import tkinter as tk
from tkinter import ttk
import matplotlib
# 配置matplotlib支持中文
matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']  # 指定默认字体
matplotlib.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题

from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from matplotlib.figure import Figure
import random
import numpy as np


class AnalyzeWindow(tk.Toplevel):
    def __init__(self, master, records):
        super().__init__(master)
        self.title("历史数据分析")
        self.geometry("1200x800")
        self.records = records
        self.colors = self._assign_colors(len(records))
        self.record_vars = [tk.BooleanVar(value=True) for _ in records]
        self.curve_vars = []  # [[BooleanVar, ...], ...]
        self.expanded = [False for _ in records]
        self.log_var = tk.BooleanVar(value=False)  # 对数显示控制
        self.legend_var = tk.BooleanVar(value=True)  # 图例显示控制
        self._build_ui()
        self._draw_plot()

    def _build_ui(self):
        main_frame = tk.Frame(self)
        main_frame.pack(fill=tk.BOTH, expand=True)
        # 左侧数据列表
        left = tk.Frame(main_frame, width=320)
        left.pack(side=tk.LEFT, fill=tk.Y, padx=10, pady=10)
        
        # 创建可滚动的框架
        canvas = tk.Canvas(left, width=300)
        scrollbar = ttk.Scrollbar(left, orient="vertical", command=canvas.yview)
        self.scrollable_frame = tk.Frame(canvas)
        
        self.scrollable_frame.bind(
            "<Configure>",
            lambda e: canvas.configure(scrollregion=canvas.bbox("all"))
        )
        
        canvas.create_window((0, 0), window=self.scrollable_frame, anchor="nw")
        canvas.configure(yscrollcommand=scrollbar.set)
        
        # 绑定滚轮事件
        def _on_mousewheel(event):
            canvas.yview_scroll(int(-1*(event.delta/120)), "units")
        
        def _bind_to_mousewheel(widget):
            widget.bind("<MouseWheel>", _on_mousewheel)
            for child in widget.winfo_children():
                _bind_to_mousewheel(child)
        
        canvas.bind("<MouseWheel>", _on_mousewheel)
        self.scrollable_frame.bind("<MouseWheel>", _on_mousewheel)
        
        # 存储绑定函数以便后续使用
        self._bind_mousewheel = _bind_to_mousewheel
        
        # 布局canvas和滚动条
        canvas.pack(side="left", fill="both", expand=True)
        scrollbar.pack(side="right", fill="y")
        
        self.listbox = []
        for i, rec in enumerate(self.records):
            color = self.colors[i]
            frame = tk.Frame(self.scrollable_frame)
            frame.pack(fill=tk.X, pady=2)
            rec_cb = tk.Checkbutton(frame, variable=self.record_vars[i], command=lambda idx=i: self._on_record_cb(idx))
            rec_cb.grid(row=0, column=0, sticky='w')
            color_label = tk.Label(frame, width=2, bg=color)
            color_label.grid(row=0, column=1, padx=4, sticky='w')
            
            # 展开按钮放在色块后面
            expand_btn = ttk.Button(frame, text="展开" if not self.expanded[i] else "收起", width=4,
                                    command=lambda idx=i: self._toggle_expand(idx))
            expand_btn.grid(row=0, column=2, padx=2, sticky='w')
            
            # 信息标签，长文字显示省略号
            info_text = f"{rec['time']}  {rec.get('remark', '')}"
            info = tk.Label(frame, text=info_text, anchor='w')
            info.grid(row=0, column=3, sticky='we', padx=(5, 0))
            
            # 设置列权重，让信息标签可伸缩
            frame.grid_columnconfigure(3, weight=1)
            self.listbox.append((rec_cb, color_label, info, expand_btn, frame))
            # 曲线checkbox
            curve_vars = []
            curve_frames = []
            data = rec.get('data', {})
            channel_order = rec.get('channel_order', None)
            
            # 检查数据格式，确保有对数数据
            if isinstance(data, dict) and 'original' in data and 'log' in data:
                # 新格式：包含原始数据和对数数据
                original_data = data['original']
                log_data = data['log']
            else:
                # 旧格式：只有一组数据，假设为原始数据
                original_data = data
                log_data = {}
            
            # 检查是否为新的数据格式（包含遍数信息）
            def is_new_format(data):
                if not data:
                    return False
                # 检查第一个通道的数据是否为字典格式（包含遍数）
                first_ch_data = next(iter(data.values()), None)
                return isinstance(first_ch_data, dict)
            
            original_is_new = is_new_format(original_data)

            if channel_order:
                # 如果有保存的通道顺序，按顺序显示
                channel_count = len([ch for ch in channel_order if ch > 0])
                for j in range(channel_count):
                    # 从channel_order中提取通道号（低4位）
                    logical_ch = (channel_order[j] & 0x0F) if j < len(channel_order) and channel_order[j] > 0 else j + 1
                    physical_ch = str(logical_ch)  # 使用逻辑通道号作为数据键

                    if (physical_ch in original_data) or (physical_ch in log_data):
                        if original_is_new and physical_ch in original_data:
                            # 新格式：为每一遍创建checkbox
                            rounds_data = original_data[physical_ch]
                            for round_no in sorted(rounds_data.keys()):
                                var = tk.BooleanVar(value=True)
                                curve_vars.append(var)
                                curve_frame = tk.Frame(self.scrollable_frame)
                                # 缩进显示
                                curve_frame.pack_forget()  # 初始不显示
                                curve_label = tk.Label(curve_frame, text=f"S{logical_ch}-第{round_no}遍", width=12)
                                curve_label.pack(side=tk.LEFT, padx=(30, 0))
                                cb = tk.Checkbutton(curve_frame, variable=var,
                                                    command=lambda ridx=i, cidx=len(curve_vars)-1: self._on_curve_cb(ridx, cidx))
                                cb.pack(side=tk.LEFT)
                                curve_frames.append(curve_frame)
                        else:
                            # 旧格式：单个checkbox
                            var = tk.BooleanVar(value=True)
                            curve_vars.append(var)
                            curve_frame = tk.Frame(self.scrollable_frame)
                            # 缩进显示
                            curve_frame.pack_forget()  # 初始不显示
                            curve_label = tk.Label(curve_frame, text=f"S{logical_ch}", width=4)
                            curve_label.pack(side=tk.LEFT, padx=(30, 0))
                            cb = tk.Checkbutton(curve_frame, variable=var,
                                                command=lambda ridx=i, cidx=j: self._on_curve_cb(ridx, cidx))
                            cb.pack(side=tk.LEFT)
                            curve_frames.append(curve_frame)
            else:
                # 兼容旧数据，按默认顺序显示
                for j, ch in enumerate(sorted(original_data.keys())):
                    if original_is_new:
                        # 新格式：为每一遍创建checkbox
                        rounds_data = original_data[ch]
                        for round_no in sorted(rounds_data.keys()):
                            var = tk.BooleanVar(value=True)
                            curve_vars.append(var)
                            curve_frame = tk.Frame(self.scrollable_frame)
                            # 缩进显示
                            curve_frame.pack_forget()  # 初始不显示
                            curve_label = tk.Label(curve_frame, text=f"S{j + 1}-第{round_no}遍", width=12)
                            curve_label.pack(side=tk.LEFT, padx=(30, 0))
                            cb = tk.Checkbutton(curve_frame, variable=var,
                                                command=lambda ridx=i, cidx=len(curve_vars)-1: self._on_curve_cb(ridx, cidx))
                            cb.pack(side=tk.LEFT)
                            curve_frames.append(curve_frame)
                    else:
                        # 旧格式：单个checkbox
                        var = tk.BooleanVar(value=True)
                        curve_vars.append(var)
                        curve_frame = tk.Frame(self.scrollable_frame)
                        # 缩进显示
                        curve_frame.pack_forget()  # 初始不显示
                        curve_label = tk.Label(curve_frame, text=f"S{j + 1}", width=4)
                        curve_label.pack(side=tk.LEFT, padx=(30, 0))
                        cb = tk.Checkbutton(curve_frame, variable=var,
                                            command=lambda ridx=i, cidx=j: self._on_curve_cb(ridx, cidx))
                        cb.pack(side=tk.LEFT)
                        curve_frames.append(curve_frame)
            self.curve_vars.append(curve_vars)
            self.listbox[-1] += (curve_frames,)
        self._update_curve_visibility()
        
        # 为所有子组件绑定滚轮事件
        self._bind_mousewheel(self.scrollable_frame)
        # 右侧大折线图
        right = tk.Frame(main_frame)
        right.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=10, pady=10)

        # 右上角控制选项
        control_frame = tk.Frame(right)
        control_frame.pack(fill=tk.X, pady=(0, 5))

        # 对数checkbox
        log_cb = tk.Checkbutton(control_frame, text="对数", variable=self.log_var, command=self._on_log_change)
        log_cb.pack(side=tk.RIGHT, padx=(5, 0))

        # 显示图例checkbox
        legend_cb = tk.Checkbutton(control_frame, text="显示图例", variable=self.legend_var,
                                   command=self._on_legend_change)
        legend_cb.pack(side=tk.RIGHT, padx=(5, 0))

        self.fig = Figure(figsize=(7, 5), dpi=100)
        self.ax = self.fig.add_subplot(111)
        self.canvas = FigureCanvasTkAgg(self.fig, master=right)
        self.canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)

        # 底部标准曲线选择与分析结果
        bottom = tk.Frame(self)
        bottom.pack(fill=tk.X, padx=10, pady=5)
        tk.Label(bottom, text="选择标准曲线：").pack(side=tk.LEFT)
        self.std_rec_var = tk.StringVar()
        self.std_curve_var = tk.StringVar()
        self.std_rec_combo = ttk.Combobox(bottom, textvariable=self.std_rec_var, state='readonly', width=20)
        self.std_curve_combo = ttk.Combobox(bottom, textvariable=self.std_curve_var, state='readonly', width=8)
        self.std_rec_combo.pack(side=tk.LEFT, padx=2)
        self.std_curve_combo.pack(side=tk.LEFT, padx=2)
        self.std_rec_combo['values'] = [f"{rec['time']}|{rec.get('remark', '')}" for rec in self.records]
        self.std_rec_combo.bind('<<ComboboxSelected>>', self._on_std_rec_change)
        self.std_curve_combo['values'] = []
        # 横轴范围选择
        tk.Label(bottom, text="横轴范围：").pack(side=tk.LEFT, padx=(10, 0))
        self.range_start_var = tk.StringVar()
        self.range_end_var = tk.StringVar()
        self.range_start_entry = ttk.Entry(bottom, textvariable=self.range_start_var, width=6)
        self.range_end_entry = ttk.Entry(bottom, textvariable=self.range_end_var, width=6)
        self.range_start_entry.pack(side=tk.LEFT, padx=2)
        tk.Label(bottom, text="~").pack(side=tk.LEFT)
        self.range_end_entry.pack(side=tk.LEFT, padx=2)
        range_btn = ttk.Button(bottom, text="应用", command=self._apply_range)
        range_btn.pack(side=tk.LEFT, padx=2)
        analyze_btn = ttk.Button(bottom, text="确定", command=self._analyze_diff)
        analyze_btn.pack(side=tk.LEFT, padx=5)
        # 结果表格
        self.result_tree = ttk.Treeview(self, columns=("color", "rec", "curve", "mse", "l1", "corr"), show="headings",
                                        height=5)
        self.result_tree.heading("color", text="颜色")
        self.result_tree.heading("rec", text="记录")
        self.result_tree.heading("curve", text="曲线")
        self.result_tree.heading("mse", text="MSE")
        self.result_tree.heading("l1", text="L1")
        self.result_tree.heading("corr", text="相关系数")
        self.result_tree.column("color", width=40, anchor='center')
        self.result_tree.column("rec", width=180)
        self.result_tree.column("curve", width=60)
        self.result_tree.column("mse", width=100)
        self.result_tree.column("l1", width=100)
        self.result_tree.column("corr", width=100)
        self.result_tree.pack(fill=tk.X, padx=10, pady=5)
        # 默认范围
        self._set_default_range()

    def _assign_colors(self, n):
        # 使用matplotlib调色板tab10，保证唯一
        import matplotlib
        tab = list(matplotlib.colormaps['tab10'].colors)
        random.shuffle(tab)
        colors = [matplotlib.colors.to_hex(tab[i % 10]) for i in range(n)]
        return colors

    def _toggle_expand(self, idx):
        self.expanded[idx] = not self.expanded[idx]
        btn = self.listbox[idx][3]
        btn.config(text="收起" if self.expanded[idx] else "展开")
        self._update_curve_visibility()

    def _update_curve_visibility(self):
        for i, (rec_cb, color_label, info, expand_btn, frame, curve_frames) in enumerate(self.listbox):
            if self.expanded[i]:
                for cf in curve_frames:
                    cf.pack(fill=tk.X, padx=40, pady=1, after=frame)
            else:
                for cf in curve_frames:
                    cf.pack_forget()

    def _on_record_cb(self, idx):
        # 勾选/取消记录时，所有曲线同步
        val = self.record_vars[idx].get()
        for var in self.curve_vars[idx]:
            var.set(val)
        self._draw_plot()

    def _on_curve_cb(self, ridx, cidx):
        # 曲线checkbox变化时，外部checkbox同步
        all_vars = self.curve_vars[ridx]
        all_checked = all(v.get() for v in all_vars)
        any_checked = any(v.get() for v in all_vars)
        self.record_vars[ridx].set(all_checked)
        self._draw_plot()

    def _on_log_change(self):
        """对数checkbox变化时重新绘图"""
        self._draw_plot()

    def _on_legend_change(self):
        """图例checkbox变化时重新绘图"""
        self._draw_plot()

    def _draw_plot(self):
        self.ax.clear()
        for i, rec in enumerate(self.records):
            color = self.colors[i]
            data = rec.get('data', {})
            channel_order = rec.get('channel_order', None)

            # 处理数据格式，确保有对数数据
            if isinstance(data, dict) and 'original' in data and 'log' in data:
                # 新格式：包含原始数据和对数数据
                original_data = data['original']
                log_data = data['log']
            else:
                # 旧格式：只有一组数据，假设为原始数据
                original_data = data
                log_data = {}
                # 为旧数据计算对数值
                import math
                for ch, vals in original_data.items():
                    log_vals = []
                    for val in vals:
                        try:
                            log_val = math.log10(abs(val)) if val != 0 else float('-inf')
                        except Exception:
                            log_val = float('-inf')
                        log_vals.append(log_val)
                    log_data[ch] = log_vals

            # 根据对数checkbox选择使用的数据
            plot_data = log_data if self.log_var.get() else original_data
            
            # 检查是否为新的数据格式（包含遍数信息）
            def is_new_format(data):
                if not data:
                    return False
                # 检查第一个通道的数据是否为字典格式（包含遍数）
                first_ch_data = next(iter(data.values()), None)
                return isinstance(first_ch_data, dict)
            
            plot_is_new = is_new_format(plot_data)
            
            # 图例标签：优先显示备注，备注为空显示时间
            label_text = rec.get('remark', '') or rec['time']

            if channel_order:
                # 如果有保存的通道顺序，按顺序显示
                channel_count = len([ch for ch in channel_order if ch > 0])
                curve_idx = 0  # 曲线索引计数器
                for j in range(channel_count):
                    # 从channel_order中提取通道号（低4位）
                    logical_ch = (channel_order[j] & 0x0F) if j < len(channel_order) and channel_order[j] > 0 else j + 1
                    physical_ch = str(logical_ch)  # 使用逻辑通道号作为数据键

                    if physical_ch in plot_data:
                        if plot_is_new:
                            # 新格式：按遍数分组
                            rounds_data = plot_data[physical_ch]
                            for round_no in sorted(rounds_data.keys()):
                                if curve_idx < len(self.curve_vars[i]) and self.curve_vars[i][curve_idx].get():
                                    vals = rounds_data[round_no]
                                    if vals:
                                        x = range(len(vals))
                                        # 为不同遍数使用不同的线型
                                        round_no_int = int(round_no)  # 确保round_no是整数
                                        line_style = '-' if round_no_int <= 5 else '--'
                                        self.ax.plot(x, vals, 
                                                   label=f"{label_text} S{logical_ch}-第{round_no}遍", 
                                                   color=color, 
                                                   linestyle=line_style)
                                curve_idx += 1
                        else:
                            # 旧格式：直接绘制
                            if curve_idx < len(self.curve_vars[i]) and self.curve_vars[i][curve_idx].get():
                                vals = plot_data[physical_ch]
                                x = range(len(vals))
                                self.ax.plot(x, vals, label=f"{label_text} S{logical_ch}", color=color)
                            curve_idx += 1
            else:
                # 兼容旧数据，按默认顺序显示
                curve_idx = 0  # 曲线索引计数器
                for j, ch in enumerate(sorted(plot_data.keys())):
                    if plot_is_new:
                        # 新格式：按遍数分组
                        rounds_data = plot_data[ch]
                        for round_no in sorted(rounds_data.keys()):
                            if curve_idx < len(self.curve_vars[i]) and self.curve_vars[i][curve_idx].get():
                                vals = rounds_data[round_no]
                                if vals:
                                    x = range(len(vals))
                                    # 为不同遍数使用不同的线型
                                    round_no_int = int(round_no)  # 确保round_no是整数
                                    line_style = '-' if round_no_int <= 5 else '--'
                                    self.ax.plot(x, vals, 
                                               label=f"{label_text} S{ch}-第{round_no}遍", 
                                               color=color, 
                                               linestyle=line_style)
                            curve_idx += 1
                    else:
                        # 旧格式：直接绘制
                        if curve_idx < len(self.curve_vars[i]) and self.curve_vars[i][curve_idx].get():
                            vals = plot_data[ch]
                            x = range(len(vals))
                            self.ax.plot(x, vals, label=f"{label_text} S{ch}", color=color)
                        curve_idx += 1

        # 画范围竖线
        if hasattr(self, 'range_start') and hasattr(self, 'range_end'):
            self.ax.axvline(self.range_start, color='black', linestyle='--')
            self.ax.axvline(self.range_end, color='black', linestyle='--')

        # 图例显示控制
        if self.legend_var.get():
            # 检查是否有带标签的线条
            lines_with_labels = [line for line in self.ax.get_lines() if line.get_label() and not line.get_label().startswith('_')]
            if lines_with_labels:
                self.ax.legend()

        self.ax.set_xlabel("采样点")
        # 根据对数checkbox设置y轴标签
        if self.log_var.get():
            self.ax.set_ylabel("对数电流 log10(I)")
        else:
            self.ax.set_ylabel("电流 I")
        self.canvas.draw()

    def _on_std_rec_change(self, event=None):
        idx = self.std_rec_combo.current()
        if idx < 0:
            self.std_curve_combo['values'] = []
            return
        rec = self.records[idx]
        data = rec.get('data', {})
        channel_order = rec.get('channel_order', None)

        if channel_order:
            # 如果有保存的通道顺序，按顺序生成S1/S2等名称
            channel_count = len([ch for ch in channel_order if ch > 0])
            curve_names = []
            for i in range(channel_count):
                physical_ch = str(i + 1)  # 物理通道号（字符串形式）
                # 从channel_order中提取通道号（低4位）
                logical_ch = (channel_order[i] & 0x0F) if i < len(channel_order) and channel_order[i] > 0 else i + 1
                if physical_ch in data:
                    curve_names.append(f"S{logical_ch}")
        else:
            # 兼容旧数据，按默认顺序生成
            curve_names = [f"S{j + 1}" for j in range(len(data))]

        self.std_curve_combo['values'] = curve_names
        if curve_names:
            self.std_curve_combo.current(0)

    def _set_default_range(self):
        # 默认范围为所有点
        maxlen = 0
        for rec in self.records:
            for vals in rec.get('data', {}).values():
                maxlen = max(maxlen, len(vals))
        self.range_start_var.set('0')
        self.range_end_var.set(str(maxlen - 1 if maxlen > 0 else 0))
        self.range_start = 0
        self.range_end = maxlen - 1 if maxlen > 0 else 0

    def _apply_range(self):
        try:
            start = int(self.range_start_var.get())
            end = int(self.range_end_var.get())
            if start < 0 or end < start:
                raise ValueError
            self.range_start = start
            self.range_end = end
        except Exception:
            self._set_default_range()
        self._draw_plot()
        self._analyze_diff()

    def _analyze_diff(self):
        # 获取标准曲线
        rec_idx = self.std_rec_combo.current()
        curve_idx = self.std_curve_combo.current()
        if rec_idx < 0 or curve_idx < 0:
            return
        std_rec = self.records[rec_idx]
        std_data = std_rec.get('data', {})
        std_keys = sorted(std_data.keys())
        if curve_idx >= len(std_keys):
            return
        std_curve = std_data[std_keys[curve_idx]]
        # 范围
        start = getattr(self, 'range_start', 0)
        end = getattr(self, 'range_end', len(std_curve) - 1)
        # 遍历所有勾选曲线
        results = []
        for i, rec in enumerate(self.records):
            color = self.colors[i]
            data = rec.get('data', {})
            for j, ch in enumerate(sorted(data.keys())):
                if self.curve_vars[i][j].get():
                    vals = data[ch]
                    # 补0对齐
                    maxlen = max(len(std_curve), len(vals))
                    std_vals = std_curve + [0] * (maxlen - len(std_curve))
                    vals2 = vals + [0] * (maxlen - len(vals))
                    # 截取范围
                    s = max(0, start)
                    e = min(end, maxlen - 1)
                    arr1 = np.array(std_vals[s:e + 1])
                    arr2 = np.array(vals2[s:e + 1])
                    # MSE
                    mse = np.mean((arr2 - arr1) ** 2)
                    # L1
                    l1 = np.sum(np.abs(arr2 - arr1))
                    # 相关系数
                    if np.std(arr1) == 0 or np.std(arr2) == 0:
                        corr = 'NaN'
                    else:
                        corr = np.corrcoef(arr1, arr2)[0, 1]
                        corr = f"{corr:.4f}"
                    results.append((color, f"{rec['time']}|{rec.get('remark', '')}", f"S{j + 1}", mse, l1, corr))
        # 显示结果
        for i in self.result_tree.get_children():
            self.result_tree.delete(i)
        for color, recname, curvename, mse, l1, corr in results:
            self.result_tree.insert('', 'end', values=('', recname, curvename, mse, l1, corr), tags=(color,))
            self.result_tree.tag_configure(color, background=color)
